
"""
构建数据分析流程图
"""
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver
from core.data_analyze.state import AnalyzeState
from core.data_analyze.nodes import (
    human_feedback_node,
    task_rewrite_node,
    fetch_data_node,
    task_decompose_node,
    data_analyze_node,
    report_node
)
from langgraph.graph.state import CompiledStateGraph
import matplotlib.pyplot as plt
import matplotlib.image as mpimg

def _build_base_graph():
    """Build and return the base state graph with all nodes and edges."""
    builder = StateGraph(AnalyzeState)
    builder.add_edge(START, 'task_decompose')
    builder.add_node("human_feedback", human_feedback_node)
    builder.add_node('task_rewrite', task_rewrite_node)
    builder.add_node("fetch_data", fetch_data_node)
    builder.add_node("task_decompose", task_decompose_node)
    builder.add_node("data_analyze", data_analyze_node)
    builder.add_node("report", report_node)
    builder.add_edge("report", END)
    return builder


def build_graph_with_memory():
    """Build and return the agent workflow graph with memory."""
    # use persistent memory to save conversation history
    # TODO: be compatible with SQLite / PostgreSQL
    memory = MemorySaver()

    builder = _build_base_graph()
    return builder.compile(checkpointer=memory)


def build_graph():
    """Build and return the agent workflow graph without memory."""
    builder = _build_base_graph()
    return builder.compile()


def show_graph(graph: CompiledStateGraph):
    """Show graph."""
    mermaid_code = graph.get_graph().draw_mermaid_png()
    with open("graph.jpg", "wb") as f:
        f.write(mermaid_code)

    img = mpimg.imread("graph.jpg")
    plt.imshow(img)
    plt.axis('off')
    plt.show()